#include <ntddk.h>

#include "../SysMonCommon/SysMonCommon.h"
#include "SysMonDriver.h"

Globals g_Globals;

VOID
SysMonUnload(_In_ PDRIVER_OBJECT DriverObject);

NTSTATUS
SysMonCreateAndClose(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp);

NTSTATUS
SysMonRed(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp);

VOID
OnProcessNotify(_In_ PEPROCESS Process, _In_ HANDLE ProcessId, _In_ PPS_CREATE_NOTIFY_INFO CreateInfo);

VOID
OnThreadNotify(_In_ HANDLE ProcessId, _In_ HANDLE ThreadId, _In_ BOOLEAN Create);

VOID
OnImageLoadNotify(_In_opt_ PUNICODE_STRING FullImageName, _In_ HANDLE ProcessId, _In_ PIMAGE_INFO ImageInfo);

NTSTATUS
OnRegistryNotify(_In_ PVOID CallbackContext, _In_opt_ PVOID Argument1, _In_opt_ PVOID Argument2);

VOID
PushItem(_In_ LIST_ENTRY* Entry);

EXTERN_C
NTSTATUS
DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegisterPath)
{
	NTSTATUS Status = STATUS_SUCCESS;

	InitializeListHead(&g_Globals.m_ItemHead);
	g_Globals.m_Mutex.Init();

	PDEVICE_OBJECT DeviceObject = nullptr;
	UNICODE_STRING SymLink      = RTL_CONSTANT_STRING(L"\\??\\SysMon");
	BOOLEAN SymLinkCreated      = FALSE;

	do 
	{
		UNICODE_STRING DevName = RTL_CONSTANT_STRING(L"\\Device\\SysMon");
		Status = IoCreateDevice(DriverObject, 0, &DevName, FILE_DEVICE_UNKNOWN, 0, TRUE, &DeviceObject);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to cxreate device (0x%08X)\n", Status));
			break;
		}
		DeviceObject->Flags |= DO_DIRECT_IO;

		Status = IoCreateSymbolicLink(&SymLink, &DevName);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "falied to create sym link (0x%08X)\n", Status));
			break;
		}
		SymLinkCreated = TRUE;

		Status = PsSetCreateProcessNotifyRoutineEx(OnProcessNotify, FALSE);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to register process callback (0x%08X)\n", Status));
			break;
		}

		Status = PsSetCreateThreadNotifyRoutine(OnThreadNotify);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to register thread callback (0x%08X)\n", Status));
			break;
		}

		Status = PsSetLoadImageNotifyRoutine(OnImageLoadNotify);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to set image load callback (status = %08X)\n", Status));
			break;
		}

		UNICODE_STRING Altitude = RTL_CONSTANT_STRING(L"1245.314");
		Status = CmRegisterCallbackEx(OnRegistryNotify, &Altitude, DriverObject, nullptr, &g_Globals.m_RegCookie, nullptr);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to set Registry callback (%08X)\n", Status));
			break;
		}

		KdPrint(("SysMonDriver loaded successfully\n"));
	} while (FALSE);

	if (NT_SUCCESS(Status) == FALSE) {
		if (SymLinkCreated == TRUE) {
			IoDeleteSymbolicLink(&SymLink);
		}
		if (DeviceObject != nullptr) {
			IoDeleteDevice(DeviceObject);
		}
		KdPrint(("SysMonDriver loading failed\n"));
	}

	DriverObject->MajorFunction[IRP_MJ_CREATE] = SysMonCreateAndClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE]  = SysMonCreateAndClose;
	DriverObject->MajorFunction[IRP_MJ_READ]   = SysMonRed;

	DriverObject->DriverUnload                 = SysMonUnload;

	return Status;
}

VOID
SysMonUnload(_In_ PDRIVER_OBJECT DriverObject)
{
	PsSetCreateProcessNotifyRoutineEx(OnProcessNotify, TRUE);
	PsRemoveCreateThreadNotifyRoutine(OnThreadNotify);
	PsRemoveLoadImageNotifyRoutine(OnImageLoadNotify);
	CmUnRegisterCallback(g_Globals.m_RegCookie);

	UNICODE_STRING SymLink = RTL_CONSTANT_STRING(L"\\??\\SysMon");
	IoDeleteSymbolicLink(&SymLink);
	IoDeleteDevice(DriverObject->DeviceObject);

	while (IsListEmpty(&g_Globals.m_ItemHead) == FALSE) {
		LIST_ENTRY* Entry = RemoveHeadList(&g_Globals.m_ItemHead);
		ExFreePool(CONTAINING_RECORD(Entry, FullItem<ItemHeader>, m_Entry));
	}

	KdPrint(("SysMonDriver uninstalled successfully\n"));
}

NTSTATUS
SysMonCreateAndClose(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp)
{
	NTSTATUS Status           = STATUS_SUCCESS;

	Irp->IoStatus.Status      = Status;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return Status;
}

NTSTATUS
SysMonRed(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp)
{
	NTSTATUS Status          = STATUS_SUCCESS;
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	ULONG Len                = Stack->Parameters.Read.Length;
	ULONG Count              = 0;

	PUCHAR Buffer = (PUCHAR)MmGetSystemAddressForMdlSafe(Irp->MdlAddress, NormalPagePriority);
	if (Buffer == nullptr) {
		Status = STATUS_INSUFFICIENT_RESOURCES;
	} else {
		AutoLock<FastMutex> Lock(g_Globals.m_Mutex);
		while (TRUE) {
			if (IsListEmpty(&g_Globals.m_ItemHead) == TRUE) {
				break;
			}

			LIST_ENTRY* Entry           = RemoveHeadList(&g_Globals.m_ItemHead);
			FullItem<ItemHeader>* Info  = CONTAINING_RECORD(Entry, FullItem<ItemHeader>, m_Entry);
			USHORT Size                 = Info->m_Data.m_Size;
			if (Len < Size) {
				InsertHeadList(&g_Globals.m_ItemHead, Entry);
				break;
			}
			::memcpy(Buffer, &Info->m_Data, Size);

			g_Globals.m_ItemCount--;
			Len    -= Size;
			Buffer += Size;
			Count  += Size;
			ExFreePool(Info);
		}
	}

	Irp->IoStatus.Status      = Status;
	Irp->IoStatus.Information = Count;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Status;
}

VOID
OnProcessNotify(_In_ PEPROCESS Process, _In_ HANDLE ProcessId, _In_ PPS_CREATE_NOTIFY_INFO CreateInfo)
{
	if (CreateInfo != nullptr) {
		USHORT AllocSize       = sizeof(FullItem<ProcessCreateInfo>);
		USHORT CommandLineSize = 0;
		if (CreateInfo->CommandLine != nullptr) {
			CommandLineSize = CreateInfo->CommandLine->Length;
			AllocSize       = AllocSize + CommandLineSize;
		}
		FullItem<ProcessCreateInfo>* Info = (FullItem<ProcessCreateInfo>*)ExAllocatePoolWithTag(PagedPool, AllocSize, DRIVER_TAG);
		if (Info == nullptr) {
			KdPrint((DRIVER_PREFIX "failed allocation\n"));
			return;
		}

		ProcessCreateInfo& Item = Info->m_Data;
		KeQuerySystemTimePrecise(&Item.m_Time);
		Item.m_Type             = ItemType::ProcessCreate;
		Item.m_Size             = sizeof(ProcessCreateInfo) + CommandLineSize;
		Item.m_ProcessId        = HandleToULong(ProcessId);
		Item.m_ParentProcessId  = HandleToULong(CreateInfo->ParentProcessId);

		if (CommandLineSize > 0) {
			::memcpy((UCHAR*)&Item + sizeof(Item), CreateInfo->CommandLine->Buffer, CommandLineSize);
			Item.m_CommandLineLength = CommandLineSize / sizeof(WCHAR);
			Item.m_CommandLineOffset = sizeof(Item);
		} else {
			Item.m_CommandLineLength = 0;
		}

		PushItem(&Info->m_Entry);
	} else {
		FullItem<ProcessExitInfo>* Info = (FullItem<ProcessExitInfo>*)ExAllocatePoolWithTag(PagedPool, sizeof(FullItem<ProcessExitInfo>), DRIVER_TAG);
		if (Info == nullptr) {
			KdPrint((DRIVER_PREFIX "falied allocation\n"));
			return;
		}

		ProcessExitInfo& Item = Info->m_Data;
		KeQuerySystemTimePrecise(&Item.m_Time);
		Item.m_Type           = ItemType::ProcessExit;
		Item.m_ProcessId      = HandleToULong(ProcessId);
		Item.m_Size           = sizeof(ProcessExitInfo);

		PushItem(&Info->m_Entry);
	}
}

VOID
OnThreadNotify(_In_ HANDLE ProcessId, _In_ HANDLE ThreadId, _In_ BOOLEAN Create)
{
	ULONG Size = sizeof(FullItem<ThreadCreateExitInfo>);
	FullItem<ThreadCreateExitInfo>* Info = (FullItem<ThreadCreateExitInfo>*)ExAllocatePoolWithTag(PagedPool, Size, DRIVER_TAG);
	if (Info == nullptr) {
		KdPrint((DRIVER_PREFIX "Failed to allocate memory\n"));
		return;
	}
	ThreadCreateExitInfo& Item = Info->m_Data;
	KeQuerySystemTimePrecise(&Item.m_Time);
	Item.m_Size                = sizeof(Item);
	Item.m_Type                = Create ? ItemType::ThreadCreate : ItemType::ThreadExit;
	Item.m_ProcessId           = HandleToLong(ProcessId);
	Item.m_ThreadId            = HandleToLong(ThreadId);

	PushItem(&Info->m_Entry);
}

VOID
OnImageLoadNotify(_In_opt_ PUNICODE_STRING FullImageName, _In_ HANDLE ProcessId, _In_ PIMAGE_INFO ImageInfo)
{
	if (ProcessId == nullptr) {
		return;
	}

	ULONG Size = sizeof(FullItem<ImageLoadInfo>);
	FullItem<ImageLoadInfo>* Info = (FullItem<ImageLoadInfo>*)ExAllocatePoolWithTag(PagedPool, Size, DRIVER_TAG);
	if (Info == nullptr) {
		KdPrint((DRIVER_PREFIX "Failed to allocate memory\n"));
		return;
	}
	
	::memset(Info, 0, Size);

	ImageLoadInfo& Item = Info->m_Data;
	KeQuerySystemTimePrecise(&Item.m_Time);
	Item.m_Size         = sizeof(Item);
	Item.m_Type         = ItemType::ImageLoad;
	Item.m_ProcessId    = HandleToULong(ProcessId);
	Item.m_ImageSize    = ImageInfo->ImageSize;
	Item.m_LoadAddress  = ImageInfo->ImageBase;

	if (FullImageName != nullptr) {
		::memcpy(Item.m_ImageFileName, FullImageName->Buffer, FullImageName->Length);
	} else {
		::wcscpy_s(Item.m_ImageFileName, L"(unknown)");
	}

	PushItem(&Info->m_Entry);
}

NTSTATUS
OnRegistryNotify(_In_ PVOID CallbackContext, _In_opt_ PVOID Argument1, _In_opt_ PVOID Argument2)
{
	static CONST WCHAR Marchine[] = L"\\REGISTRY\\MACHINE\\";

	if (Argument2 == nullptr || Argument2 == nullptr) {
		return STATUS_SUCCESS;
	}

	switch ((REG_NOTIFY_CLASS)(ULONG_PTR)Argument1) {
	case RegNtPostSetValueKey: {
		REG_POST_OPERATION_INFORMATION* Args = static_cast<REG_POST_OPERATION_INFORMATION*>(Argument2);
		if (NT_SUCCESS(Args->Status) == FALSE) {
			break;
		}

		PCUNICODE_STRING Name;
		if (NT_SUCCESS(CmCallbackGetKeyObjectIDEx(&g_Globals.m_RegCookie, Args->Object, nullptr, &Name, 0)) != FALSE) {
			if (::wcsncmp(Name->Buffer, Marchine, ARRAYSIZE(Marchine) - 1) == 0) {
				REG_SET_VALUE_KEY_INFORMATION* PreInfo = static_cast<REG_SET_VALUE_KEY_INFORMATION*>(Args->PreInformation);
				NT_ASSERT(PreInfo);

				ULONG Size = sizeof(FullItem<RegistrySetValueInfo>);
				FullItem<RegistrySetValueInfo>* Info = static_cast<FullItem<RegistrySetValueInfo>*>(ExAllocatePoolWithTag(PagedPool, Size, DRIVER_TAG));
				if (Info == nullptr) {
					break;
				}

				RtlZeroMemory(Info, Size);
				RegistrySetValueInfo& Item = Info->m_Data;
				KeQuerySystemTimePrecise(&Item.m_Time);
				Item.m_Size = sizeof(Item);
				Item.m_Type = ItemType::RegistrySetValue;
				::wcsncpy_s(Item.m_KeyName, Name->Buffer, Name->Length / sizeof(WCHAR) - 1);
				::wcsncpy_s(Item.m_ValueName, PreInfo->ValueName->Buffer, PreInfo->ValueName->Length / sizeof(WCHAR) - 1);
				Item.m_DataType = PreInfo->Type;
				Item.m_DataSize = PreInfo->DataSize;
				Item.m_ProcessId = HandleToULong(PsGetCurrentProcessId());
				Item.m_ThreadId = HandleToULong(PsGetCurrentThreadId());
				::memcpy(Item.m_Data, PreInfo->Data, Item.m_DataSize < sizeof(Item.m_Data) ? Item.m_DataSize : sizeof(Item.m_Data));

				PushItem(&Info->m_Entry);
			}
			CmCallbackReleaseKeyObjectIDEx(Name);
		}
		break;
	}
	}

	return STATUS_SUCCESS;
}

VOID
PushItem(_In_ LIST_ENTRY* Entry)
{
	AutoLock<FastMutex> Lock(g_Globals.m_Mutex);
	if (g_Globals.m_ItemCount > 1024) {
		LIST_ENTRY* Head = RemoveHeadList(&g_Globals.m_ItemHead);
		g_Globals.m_ItemCount--;
		FullItem<ItemHeader>* Item = CONTAINING_RECORD(Head, FullItem<ItemHeader>, m_Entry);
		ExFreePool(Item);
	}
	InsertHeadList(&g_Globals.m_ItemHead, Entry);
	g_Globals.m_ItemCount++;
}